#include <ntddk.h>
#include <tdikrnl.h>
#include "sock.h"

 WCHAR *TCP_DEVICE = L"\\Device\\Tcp";
 WCHAR *UDP_DEVICE = L"\\Device\\Udp";
 WCHAR *RAWIP_DEVICE = L"\\Device\\RawIp";
PDEVICE_OBJECT
g_tcpfltobj = nullptr,		// \Device\Tcp
g_udpfltobj = nullptr,		// \Device\Udp
g_ipfltobj = nullptr,		// \Device\RawIp

g_tcpoldobj = nullptr,
g_udpoldobj = nullptr,
g_ipoldobj = nullptr,
g_devcontrol = nullptr,	// control device (exclusive access only!)
g_devnfo = nullptr;		// information device 



BOOLEAN g_got_log = FALSE;	// got log app
NTSTATUS c_n_a_device(PDRIVER_OBJECT DriverObject,
	PDEVICE_OBJECT *fltobj,
	PDEVICE_OBJECT *oldobj,
	WCHAR* devname)
{
	NTSTATUS status;
	UNICODE_STRING str;

	//
	status = IoCreateDevice(DriverObject,
		0,
		nullptr,
		FILE_DEVICE_UNKNOWN,
		0, TRUE,
		fltobj);
	if (STATUS_SUCCESS != status) {
		KdPrint(("[tdi_test] c_n_a_device:IoCreateDevice(%S):0x&x\n", devname, status));
		return status;
	}
	
	(*fltobj)->Flags |= DO_DIRECT_IO;

	RtlInitUnicodeString(&str, devname);
	status = IoAttachDevice(*fltobj, &str, oldobj);
	if (status != STATUS_SUCCESS) {
		KdPrint(("[tdi_test] c_n_a_device:IoAttachDevice(%S):0x&x\n", devname, status));
		return status;
	}
}

VOID OnUnload(IN PDRIVER_OBJECT DriverObject)
{
	UNREFERENCED_PARAMETER(DriverObject);
	if (g_tcpoldobj != nullptr) IoDetachDevice(g_tcpoldobj);
	if (g_tcpfltobj != nullptr) IoDetachDevice(g_tcpfltobj);
	if (g_updoldobj != nullptr) IoDetachDevice(g_updoldobj);
	if (g_udpfltobj != nullptr) IoDetachDevice(g_udpfltobj);
	if (g_ipoldobj != nullptr) IoDetachDevice(g_ipoldobj);
	if (g_ipfltobj != nullptr) IoDetachDevice(g_ipfltobj);
}

PDEVICE_OBJECT get_original_devobj(PDEVICE_OBJECT flt_devobj, int *proto)
{
	PDEVICE_OBJECT result;
	int ipproto;

	if (flt_devobj == g_tcpfltobj) {
		result = g_tcpfltobj;
		ipproto = IPPROTO_TCP;
	}
	else if (flt_devobj == g_udpfltobj) {
		result = g_udpfltobj;
		ipproto = IPPROTO_UDP;
	}
	else if (flt_devobj = g_ipfltobj) {
		result = g_ipfltobj;
		ipproto = IPPROTO_IP;
	}
	else {
		KdPrint(("[tdi_test] get_original_devobj:Unkown DeviceObject 0x&x\n", flt_devobj));
		ipproto = IPPROTO_IP;
		result = nullptr;
	}

	if (result != nullptr && proto != nullptr)
		*proto = ipproto;
	return result;
}
NTSTATUS DeviceDispatch(IN PDEVICE_OBJECT DeviceObject, IN PIRP irp)
{
	NTSTATUS status;
	PDEVICE_OBJECT old_devobj = get_original_devobj(DeviceObject, nullptr);
	if (nullptr != old_devobj){
		IoSkipCurrentIrpStackLocation(irp);
		status = IoCallDriver(old_devobj, irp);
	}
	else {
		status = irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
		IoCompleteRequest(irp, IO_NO_INCREMENT);
	}
	return status;
}


NTSTATUS
DriverEntry(IN PDRIVER_OBJECT theDriverObject,
	IN PUNICODE_STRING theRegistryPath)
{
	UNREFERENCED_PARAMETER(theRegistryPath);
	NTSTATUS status = STATUS_SUCCESS;
	int i;
	for each(auto EachFunc in theDriverObject->MajorFunction)
	{
		EachFunc = DeviceDispatch;
	}
	theDriverObject->DriverUnload = OnUnload;

	status = c_n_a_device(theDriverObject, &g_tcpfltobj, &g_tcpoldobj, TCP_DEVICE);
	if (status != STATUS_SUCCESS) {
		KdPrint(("[tdi_test] DriverEntry:c_n_a_device call for tcp 0x&x\n", status));
		goto done;
	}

	status = c_n_a_device(theDriverObject, &g_udpfltobj, &g_udpoldobj, UDP_DEVICE);
	if (status != STATUS_SUCCESS) {
		KdPrint(("[tdi_test] DriverEntry:c_n_a_device call for udp 0x&x\n", status));
		goto done;
	}

	status = c_n_a_device(theDriverObject, &g_ipfltobj, &g_ipoldobj, RAWIP_DEVICE);
	if (status != STATUS_SUCCESS) {
		KdPrint(("[tdi_test] DriverEntry:c_n_a_device call for rawip 0x&x\n", status));
		goto done;
	}
done:
	if (status != STATUS_SUCCESS) {
		OnUnload(theDriverObject);
	}
	return status;
}